﻿

#include "utilityCommunicator.hpp"

#include "utilityContainer.hpp"

namespace UTILITY {

void Communicator::registerTask(TaskName task) {
  unordered_map<TaskName, int>::iterator iter = taskTags_.find(task);
  if (iter == taskTags_.end()) {
    taskTags_.insert(pair<TaskName, int>(task, taskCount_++));
  }
}

MPI_Request *Communicator::insertRequest(TaskName task) {
  TaskRequestTable::iterator iter = taskRequests_.find(task);
  if (iter != taskRequests_.end()) {
    MPI_Request req;
    iter->second.push_back(req);
  } else {
    MPI_Request req;
    taskRequests_.insert(
        TaskRequestPair(task, std::vector<MPI_Request>(1, req)));
    iter = taskRequests_.find(task);

    messageNum_++;
  }
  registerTask(task);

  return &(iter->second.back());
}

int BinomialCoefficient(const int n, const int k) {
  std::vector<int> aSolutions(k);
  aSolutions[0] = n - k + 1;

  for (int i = 1; i < k; ++i) {
    aSolutions[i] = aSolutions[i - 1] * (n - k + 1 + i) / (i + 1);
  }

  return aSolutions[k - 1];
}

int Communicator::generateTag(const TaskName &task, const int source,
                              const int dest) {
  unordered_map<TaskName, int>::iterator iter = taskTags_.find(task);
  int taskNum;
  if (iter != taskTags_.end()) {
    taskNum = iter->second;
  } else {
    taskNum = taskCount_;
    taskTags_.insert(pair<TaskName, int>(task, taskCount_++));
  }

  return source >= dest ? (source * source) * (1 + taskNum) + source + dest
                        : (dest * dest) * (1 + taskNum) + source;
}

int Communicator::unregisterTask(TaskName task) {
  TaskRequestTable::iterator iter = taskRequests_.find(task);
  if (iter != taskRequests_.end()) taskRequests_.erase(task);
  return 0;
}

int Communicator::bcast(const TaskName task, void *buffer, int count,
                        int root) {
#ifdef MPIv3
  MPI_Ibcast(buffer, count, MPI_BYTE, root, this->comm_, insertRequest(task));
#else
  MPI_Bcast(buffer, count, MPI_BYTE, root, this->comm_);
#endif
  return 0;
}

int Communicator::gather(const TaskName task, const void *sendbuf,
                         int sendcount, void *recvbuf, int recvcount,
                         int root) {
#ifdef MPIv3
  MPI_Igather(sendbuf, sendcount, MPI_BYTE, recvbuf, recvcount, MPI_BYTE, root,
              this->comm_, insertRequest(task));
#else
  MPI_Gather(sendbuf, sendcount, MPI_BYTE, recvbuf, recvcount, MPI_BYTE, root,
             this->comm_);
#endif
  return 0;
}

int Communicator::gatherV(const TaskName task, const void *sendbuf,
                          int sendcount, void *recvbuf, const int *recvcounts,
                          int root) {
  int size = this->getMySize();
  int *displs = new int[size];
  displs[0] = 0;
  for (int i = 1; i < size; i++) displs[i] = displs[i - 1] + recvcounts[i - 1];
#ifdef MPIv3
  MPI_Igatherv(sendbuf, sendcount, MPI_BYTE, recvbuf, recvcounts, displs,
               MPI_BYTE, root, this->comm_, insertRequest(task));
#else
  MPI_Gatherv(sendbuf, sendcount, MPI_BYTE, recvbuf, recvcounts, displs,
              MPI_BYTE, root, this->comm_);
#endif
  delete[] displs;
  return 0;
}

int Communicator::allGather(const TaskName task, const void *sendbuf,
                            int sendcount, void *recvbuf, int recvcount) {
#ifdef MPIv3
  MPI_Iallgather(sendbuf, sendcount, MPI_BYTE, recvbuf, recvcount, MPI_BYTE,
                 this->comm_, insertRequest(task));
#else
  MPI_Allgather(sendbuf, sendcount, MPI_BYTE, recvbuf, recvcount, MPI_BYTE,
                this->comm_);
#endif
  return 0;
}

int Communicator::allGatherV(const TaskName task, const void *sendbuf,
                             int sendcount, void *recvbuf,
                             const int *recvcounts) {
  int size = this->getMySize();
  int *displs = new int[size];
  displs[0] = 0;
  for (int i = 1; i < size; i++) displs[i] = displs[i - 1] + recvcounts[i - 1];
#ifdef MPIv3
  MPI_Iallgatherv(sendbuf, sendcount, MPI_BYTE, recvbuf, recvcounts, displs,
                  MPI_BYTE, this->comm_, insertRequest(task));
#endif
  MPI_Allgatherv(sendbuf, sendcount, MPI_BYTE, recvbuf, recvcounts, displs,
                 MPI_BYTE, this->comm_);
  delete[] displs;
  return 0;
}

int Communicator::allToall(const TaskName Task, const void *sendbuf,
                           int sendcount, void *recvbuff, int recvcount) {
#ifdef MPIv3
  MPI_Ialltoall(sendbuf, sendcount, MPI_BYTE, recvbuff, recvcount, MPI_BYTE,
                this->comm_, insertRequest(task));
#else
  MPI_Alltoall(sendbuf, sendcount, MPI_BYTE, recvbuff, recvcount, MPI_BYTE,
               this->comm_);
#endif
  return 0;
}

int Communicator::allToallv(const TaskName Task, const void *sendbuf,
                            int *sendcounts, void *recvbuff, int *recvcounts) {
  int size = this->getMySize();
  int *sendDispls = new int[size];
  int *recvDispls = new int[size];
  sendDispls[0] = 0;
  recvDispls[0] = 0;
  for (int i = 1; i < size; i++) {
    sendDispls[i] = sendDispls[i - 1] + sendcounts[i - 1];
    recvDispls[i] = recvDispls[i - 1] + recvcounts[i - 1];
  }

#ifdef MPIv3
  MPI_Ialltoallv(sendbuf, sendcounts, sendDispls, MPI_BYTE, recvbuff,
                 recvcounts, recvDispls, MPI_BYTE, this->comm_,
                 insertRequest(task));
#else
  MPI_Alltoallv(sendbuf, sendcounts, sendDispls, MPI_BYTE, recvbuff, recvcounts,
                recvDispls, MPI_BYTE, this->comm_);
#endif
  delete[] sendDispls;
  delete[] recvDispls;
  return 0;
}

int Communicator::scatter(const TaskName task, const void *sendbuf,
                          int sendcount, void *recvbuf, int recvcount,
                          int root) {
#ifdef MPIv3
  MPI_Iscatter(sendbuf, sendcount, MPI_BYTE, recvbuf, recvcount, MPI_BYTE, root,
               this->comm_, insertRequest(task));
#else
  MPI_Scatter(sendbuf, sendcount, MPI_BYTE, recvbuf, recvcount, MPI_BYTE, root,
              this->comm_);
#endif
  return 0;
}

int Communicator::scatterV(const TaskName task, const void *sendbuf,
                           const int *sendcounts, void *recvbuf, int recvcount,
                           int root) {
  int size = this->getMySize();
  int *displs = new int[size];
  displs[0] = 0;
  for (int i = 1; i < size; i++) displs[i] = displs[i - 1] + sendcounts[i - 1];
#ifdef MPIv3
  MPI_Iscatterv(sendbuf, sendcounts, displs, MPI_BYTE, recvbuf, recvcount,
                MPI_BYTE, root, this->comm_, insertRequest(task));
#else
  MPI_Scatterv(sendbuf, sendcounts, displs, MPI_BYTE, recvbuf, recvcount,
               MPI_BYTE, root, this->comm_);
#endif
  delete[] displs;
  return 0;
}

int Communicator::send(const TaskName task, const void *sendbuf, int sendcount,
                       int senddest) {
  MPI_Isend(sendbuf, sendcount, MPI_BYTE, senddest,
            generateTag(task, getMyId(), senddest), this->comm_,
            insertRequest(task));
  return 0;
}

int Communicator::recv(const TaskName task, void *recvbuf, int recvcount,
                       int recvdest) {
  MPI_Irecv(recvbuf, recvcount, MPI_BYTE, recvdest,
            generateTag(task, recvdest, getMyId()), this->comm_,
            insertRequest(task));
  return 0;
}

int Communicator::groupSend(const TaskName task, const void *sendbuf,
                            int sendnum, const int *sendcounts,
                            const int *senddests) {
  size_t displ = 0;
  for (int i = 0; i < sendnum; i++) {
    MPI_Isend(&((const char *)sendbuf)[displ], sendcounts[i], MPI_BYTE,
              senddests[i], generateTag(task, getMyId(), senddests[i]),
              this->comm_, insertRequest(task));
    displ += sendcounts[i];
  }
  return 0;
}

int Communicator::groupRecv(const TaskName task, void *recvbuf, int recvnum,
                            const int *recvcounts, const int *recvdests) {
  int displ = 0;
  for (int i = 0; i < recvnum; i++) {
    MPI_Irecv(&((char *)recvbuf)[displ], recvcounts[i], MPI_BYTE, recvdests[i],
              generateTag(task, recvdests[i], this->getMyId()), this->comm_,
              insertRequest(task));
    displ += recvcounts[i];
  }
  return 0;
}

int Communicator::exchange(const TaskName task, const void *sendbuf,
                           void *recvbuf, int num, const int *counts,
                           const int *dests) {
  int displ = 0;
  for (int i = 0; i < num; i++) {
    MPI_Isend(&((const char *)sendbuf)[displ], counts[i], MPI_BYTE, dests[i],
              generateTag(task, getMyId(), dests[i]), this->comm_,
              insertRequest(task));
    MPI_Irecv(&((char *)recvbuf)[displ], counts[i], MPI_BYTE, dests[i],
              generateTag(task, dests[i], getMyId()), this->comm_,
              insertRequest(task));
    displ += counts[i];
  }
  return 0;
}

int Communicator::reduce(const TaskName task, const void *sendbuf,
                         void *recvbuf, int count, CommData datatype, CommOp op,
                         int root) {
#ifdef MPIv3
  MPI_Ireduce(sendbuf, recvbuf, count, datatype, op, root, this->comm_,
              insertRequest(task));
#else
  MPI_Reduce(sendbuf, recvbuf, count, datatype, op, root, this->comm_);
#endif
  return 0;
}

int Communicator::allReduce(const TaskName task, const void *sendbuf,
                            void *recvbuf, int count, CommData datatype,
                            CommOp op) {
#ifdef MPIv3
  MPI_Iallreduce(sendbuf, recvbuf, count, datatype, op, this->comm_,
                 insertRequest(task));
#else
  MPI_Allreduce(sendbuf, recvbuf, count, datatype, op, this->comm_);
#endif
  return 0;
}

int Communicator::finishTask(const TaskName task) {
  if (messageNum_ > COMM_MAX_MSG_NUM || messageNum_ < 0) {
    cerr << __FILE__ << " +" << __LINE__ << ":" << endl
         << __FUNCTION__ << ":" << endl
         << "Error: buffered message number exceed the limit "
         << COMM_MAX_MSG_NUM << endl;
    cerr << "Check and make sure there's no misuse of sending like interfaces"
         << endl;
    exit(-1);
  }

  TaskRequestTable::iterator iter = taskRequests_.find(task);

  if (iter != taskRequests_.end()) {
    std::vector<MPI_Request> &reqv = iter->second;
    MPI_Request *reqHandle = reqv.data();
    int count = reqv.size();
    MPI_Status *stat = new MPI_Status[count];
    MPI_Waitall(count, reqHandle, stat);
    delete[] stat;
    unregisterTask(task);
    Communicator::messageNum_--;
  }

  unordered_map<TaskName, int>::iterator tagIter = taskTags_.find(task);
  if (tagIter != taskTags_.end()) {
    taskTags_.erase(task);
    if (taskTags_.empty()) taskCount_ = 0;
  }
  return 0;
}

template <typename T>
int Communicator::exchangeArray(std::vector<T> &data,
                                const std::vector<int> &partition) {
  int arraySize = partition.size();
  std::vector<int> sendCount(size_, 0);
  for (int i = 0; i < arraySize; i++) sendCount[partition[i]]++;
  std::vector<std::vector<T> > sendBuff(size_);

  int *recvCount = new int[size_ + 1];
  allToall("collectCount", sendCount.data(), sizeof(int), &recvCount[1],
           sizeof(int));
  for (int i = 0; i < size_; i++) sendBuff[i].reserve(arraySize / size_);
  for (int i = 0; i < arraySize; i++) sendBuff[partition[i]].push_back(data[i]);
  finishTask("collectCount");
  recvCount[0] = 0;
  for (int i = 1; i <= size_; ++i) recvCount[i] += recvCount[i - 1];
  std::vector<T> dataNew(recvCount[size_]);

  for (int i = 0; i < size_; ++i) {
    if (rank_ == i) continue;
    send("exchangeData", sendBuff[i].data(), sendCount[i] * sizeof(T), i);
  }
  for (int i = 0; i < size_; ++i) {
    if (rank_ == i) {
      memcpy(&dataNew[recvCount[i]], sendBuff[i].data(),
             sendCount[i] * sizeof(T));
      continue;
    }
    recv("exchangeData", &dataNew[recvCount[i]],
         (recvCount[i + 1] - recvCount[i]) * sizeof(T), i);
  }

  finishTask("exchangeData");
  delete[] recvCount;
  data.swap(dataNew);
  return 0;
}

template <typename T>
int Communicator::exchangeArrayArray(std::vector<T> &data,
                                     std::vector<label> &offset,
                                     const std::vector<int> &partition) {
  int arraySize = partition.size();
  std::vector<int> sendCount(size_, 0), dataSendCount(size_, 0);
  for (int i = 0; i < arraySize; i++) sendCount[partition[i]]++;

  int *recvCount = new int[size_ + 1];
  allToall("collectCount", sendCount.data(), sizeof(int), &recvCount[1],
           sizeof(int));
  std::vector<int> numPerRow(arraySize);
  for (int i = 0; i < arraySize; i++) numPerRow[i] = offset[i + 1] - offset[i];

  std::vector<std::vector<int> > offsetSendBuff(size_);
  std::vector<std::vector<T> > dataSendBuff(size_);
  for (int i = 0; i < size_; i++) {
    offsetSendBuff[i].reserve(arraySize / size_);
    dataSendBuff[i].reserve(data.size() / size_);
  }
  for (int i = 0; i < arraySize; i++) {
    offsetSendBuff[partition[i]].push_back(numPerRow[i]);
    dataSendCount[partition[i]] += numPerRow[i];
    dataSendBuff[partition[i]].insert(dataSendBuff[partition[i]].end(),
                                      data.begin() + offset[i],
                                      data.begin() + offset[i + 1]);
  }

  finishTask("collectCount");
  recvCount[0] = 0;
  for (int i = 1; i <= size_; ++i) recvCount[i] += recvCount[i - 1];

  std::vector<int> offsetNew(recvCount[size_]);

  for (int i = 0; i < size_; ++i) {
    if (rank_ == i) continue;
    send("exchangeOffset", offsetSendBuff[i].data(), sendCount[i] * sizeof(int),
         i);
  }
  for (int i = 0; i < size_; ++i) {
    if (rank_ == i) {
      memcpy(&offsetNew[recvCount[i]], offsetSendBuff[i].data(),
             sendCount[i] * sizeof(int));
      continue;
    }
    recv("exchangeOffset", &offsetNew[recvCount[i]],
         (recvCount[i + 1] - recvCount[i]) * sizeof(int), i);
  }
  finishTask("exchangeOffset");

  int *dataRecvCount = new int[size_ + 1];
  dataRecvCount[0] = 0;
  for (int i = 0; i < size_; ++i) {
    dataRecvCount[i + 1] = 0;
    for (int j = recvCount[i]; j < recvCount[i + 1]; j++)
      dataRecvCount[i + 1] += offsetNew[j];
  }
  for (int i = 1; i <= size_; ++i) dataRecvCount[i] += dataRecvCount[i - 1];

  std::vector<T> dataNew(dataRecvCount[size_]);

  for (int i = 0; i < size_; ++i) {
    if (rank_ == i) continue;
    send("exchangeData", dataSendBuff[i].data(), dataSendCount[i] * sizeof(T),
         i);
  }
  for (int i = 0; i < size_; ++i) {
    if (rank_ == i) {
      memcpy(&dataNew[dataRecvCount[i]], dataSendBuff[i].data(),
             dataSendCount[i] * sizeof(T));
      continue;
    }
    recv("exchangeData", &dataNew[dataRecvCount[i]],
         (dataRecvCount[i + 1] - dataRecvCount[i]) * sizeof(T), i);
  }
  std::vector<label> offset_res(recvCount[size_] + 1);
  offset_res[0] = 0;
  for (int i = 0; i < recvCount[size_]; i++)
    offset_res[i + 1] = offset_res[i] + offsetNew[i];
  offset.swap(offset_res);
  finishTask("exchangeData");

  data.swap(dataNew);

  delete[] recvCount;
  delete[] dataRecvCount;

  return 0;
}

template int Communicator::exchangeArray(std::vector<int32_t> &data,
                                         const std::vector<int> &partition);
template int Communicator::exchangeArray(std::vector<int64_t> &data,
                                         const std::vector<int> &partition);
template int Communicator::exchangeArray(std::vector<float> &data,
                                         const std::vector<int> &partition);
template int Communicator::exchangeArray(std::vector<double> &data,
                                         const std::vector<int> &partition);
template int Communicator::exchangeArrayArray(
    std::vector<int32_t> &data, std::vector<label> &offset,
    const std::vector<int> &partition);
template int Communicator::exchangeArrayArray(
    std::vector<int64_t> &data, std::vector<label> &offset,
    const std::vector<int> &partition);
template int Communicator::exchangeArrayArray(
    std::vector<float> &data, std::vector<label> &offset,
    const std::vector<int> &partition);
template int Communicator::exchangeArrayArray(
    std::vector<double> &data, std::vector<label> &offset,
    const std::vector<int> &partition);

}  // namespace UTILITY
